import numpy as np
from scipy.linalg import eig
from typing import List, Tuple, Optional
from enum import Enum
from laminated_plate import LaminatedPlate, Layer
from fg_grc_plate import MaterialProperties

class BoundaryCondition(Enum):
    """Boundary condition types."""
    SIMPLY_SUPPORTED = "S"  # Simply supported
    CLAMPED = "C"          # Clamped
    FREE = "F"             # Free

class DifferentialQuadrature:
    """
    Differential Quadrature Method for solving plate vibration problems.
    """
    
    def __init__(self, n_points: int):
        self.n_points = n_points
        self.grid_points = None
        self.weight_matrices = None
        
    def chebyshev_gauss_lobatto_points(self) -> np.ndarray:
        """Generate Chebyshev-Gauss-Lobatto grid points."""
        i = np.arange(self.n_points)
        xi = np.cos(np.pi * i / (self.n_points - 1))
        return xi
    
    def weight_coefficients(self, xi: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate weight coefficients for first and second derivatives.
        
        Args:
            xi: Grid points
            
        Returns:
            C1, C2: Weight matrices for first and second derivatives
        """
        N = len(xi)
        C1 = np.zeros((N, N))
        C2 = np.zeros((N, N))
        
        # Calculate first derivative weights
        for i in range(N):
            for j in range(N):
                if i != j:
                    prod = 1.0
                    for k in range(N):
                        if k != i and k != j:
                            prod *= (xi[i] - xi[k])
                    
                    C1[i, j] = prod / ((xi[i] - xi[j]) * np.prod([xi[j] - xi[k] for k in range(N) if k != j]))
                    
            # Diagonal terms
            C1[i, i] = -np.sum(C1[i, :])
        
        # Calculate second derivative weights
        for i in range(N):
            for j in range(N):
                if i != j:
                    C2[i, j] = 2 * C1[i, j] * (C1[i, i] - 1/(xi[i] - xi[j]))
                else:
                    C2[i, i] = -np.sum(C2[i, :])
        
        return C1, C2
    
    def setup_grid(self):
        """Setup the DQ grid and weight matrices."""
        self.grid_points = self.chebyshev_gauss_lobatto_points()
        C1, C2 = self.weight_coefficients(self.grid_points)
        self.weight_matrices = (C1, C2)

class PlateVibrationSolver:
    """
    Solver for free vibration analysis of laminated plates using DQ method.
    """
    
    def __init__(self, plate: LaminatedPlate, n_points_x: int = 11, n_points_y: int = 11):
        self.plate = plate
        self.n_points_x = n_points_x
        self.n_points_y = n_points_y
        
        # Setup DQ grids
        self.dq_x = DifferentialQuadrature(n_points_x)
        self.dq_y = DifferentialQuadrature(n_points_y)
        self.dq_x.setup_grid()
        self.dq_y.setup_grid()
        
        # Transform grid points to physical coordinates
        self.x_points = self.plate.length * (self.dq_x.grid_points + 1) / 2
        self.y_points = self.plate.width * (self.dq_y.grid_points + 1) / 2
        
        # Create 2D coordinate meshes
        self.X, self.Y = np.meshgrid(self.x_points, self.y_points)
        
    def apply_boundary_conditions(self, K: np.ndarray, M: np.ndarray, 
                                 bc_x: Tuple[BoundaryCondition, BoundaryCondition],
                                 bc_y: Tuple[BoundaryCondition, BoundaryCondition]) -> Tuple[np.ndarray, np.ndarray]:
        """
        Apply boundary conditions to stiffness and mass matrices.
        
        Args:
            K, M: Global stiffness and mass matrices
            bc_x: Boundary conditions at x=0 and x=a
            bc_y: Boundary conditions at y=0 and y=b
            
        Returns:
            Modified K and M matrices
        """
        # Number of DOFs per node (u0, v0, w0, phi_x, phi_y)
        dof_per_node = 5
        total_nodes = self.n_points_x * self.n_points_y
        
        # Identify constrained DOFs based on boundary conditions
        constrained_dofs = []
        
        # X-direction boundaries
        for j in range(self.n_points_y):
            # x = 0 boundary
            node_idx_0 = j * self.n_points_x + 0
            # x = a boundary  
            node_idx_a = j * self.n_points_x + (self.n_points_x - 1)
            
            for node_idx, bc in [(node_idx_0, bc_x[0]), (node_idx_a, bc_x[1])]:
                base_dof = node_idx * dof_per_node
                
                if bc == BoundaryCondition.SIMPLY_SUPPORTED:
                    # w0 = 0, phi_x = 0 (for x boundaries)
                    constrained_dofs.extend([base_dof + 2, base_dof + 3])
                elif bc == BoundaryCondition.CLAMPED:
                    # w0 = 0, phi_x = 0, phi_y = 0
                    constrained_dofs.extend([base_dof + 2, base_dof + 3, base_dof + 4])
        
        # Y-direction boundaries
        for i in range(self.n_points_x):
            # y = 0 boundary
            node_idx_0 = 0 * self.n_points_x + i
            # y = b boundary
            node_idx_b = (self.n_points_y - 1) * self.n_points_x + i
            
            for node_idx, bc in [(node_idx_0, bc_y[0]), (node_idx_b, bc_y[1])]:
                base_dof = node_idx * dof_per_node
                
                if bc == BoundaryCondition.SIMPLY_SUPPORTED:
                    # w0 = 0, phi_y = 0 (for y boundaries)
                    constrained_dofs.extend([base_dof + 2, base_dof + 4])
                elif bc == BoundaryCondition.CLAMPED:
                    # w0 = 0, phi_x = 0, phi_y = 0
                    constrained_dofs.extend([base_dof + 2, base_dof + 3, base_dof + 4])
        
        # Remove duplicate DOFs
        constrained_dofs = list(set(constrained_dofs))
        
        # Create unconstrained DOF indices
        all_dofs = set(range(total_nodes * dof_per_node))
        unconstrained_dofs = sorted(list(all_dofs - set(constrained_dofs)))
        
        # Extract unconstrained portions of matrices
        K_reduced = K[np.ix_(unconstrained_dofs, unconstrained_dofs)]
        M_reduced = M[np.ix_(unconstrained_dofs, unconstrained_dofs)]
        
        return K_reduced, M_reduced
    
    def assemble_matrices(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Assemble global stiffness and mass matrices.
        
        Returns:
            K, M: Global stiffness and mass matrices
        """
        # Calculate laminate properties
        A, B, D = self.plate.calculate_laminate_stiffness()
        A_s = self.plate.calculate_shear_stiffness()
        I0, I1, I2 = self.plate.calculate_mass_inertia()
        
        # Number of DOFs
        total_nodes = self.n_points_x * self.n_points_y
        total_dofs = total_nodes * 5  # 5 DOFs per node
        
        K = np.zeros((total_dofs, total_dofs))
        M = np.zeros((total_dofs, total_dofs))
        
        # Get weight matrices
        C1_x, C2_x = self.dq_x.weight_matrices
        C1_y, C2_y = self.dq_y.weight_matrices
        
        # Scaling factors for coordinate transformation
        scale_x = 2 / self.plate.length
        scale_y = 2 / self.plate.width
        
        # Assembly loop
        for j in range(self.n_points_y):
            for i in range(self.n_points_x):
                node_idx = j * self.n_points_x + i
                base_dof = node_idx * 5
                
                # Stiffness matrix contributions
                for jj in range(self.n_points_y):
                    for ii in range(self.n_points_x):
                        node_idx_2 = jj * self.n_points_x + ii
                        base_dof_2 = node_idx_2 * 5
                        
                        # Membrane stiffness (A matrix)
                        if i == ii and j == jj:
                            # Diagonal terms
                            K[base_dof:base_dof+2, base_dof:base_dof+2] += A[:2, :2]
                        
                        # Bending stiffness (D matrix)
                        if i == ii and j == jj:
                            K[base_dof+3:base_dof+5, base_dof+3:base_dof+5] += D[:2, :2]
                        
                        # Shear stiffness
                        if i == ii and j == jj:
                            K[base_dof+2, base_dof+2] += A_s[0, 0] + A_s[1, 1]
                            K[base_dof+2, base_dof+3] += A_s[1, 1] * scale_x
                            K[base_dof+2, base_dof+4] += A_s[0, 0] * scale_y
                            K[base_dof+3, base_dof+2] += A_s[1, 1] * scale_x
                            K[base_dof+4, base_dof+2] += A_s[0, 0] * scale_y
                
                # Mass matrix contributions
                M[base_dof, base_dof] = I0     # u0
                M[base_dof+1, base_dof+1] = I0 # v0
                M[base_dof+2, base_dof+2] = I0 # w0
                M[base_dof+3, base_dof+3] = I2 # phi_x
                M[base_dof+4, base_dof+4] = I2 # phi_y
                
                # Coupling terms
                M[base_dof, base_dof+3] = M[base_dof+3, base_dof] = I1
                M[base_dof+1, base_dof+4] = M[base_dof+4, base_dof+1] = I1
        
        return K, M
    
    def solve_eigenvalue_problem(self, bc_x: Tuple[BoundaryCondition, BoundaryCondition],
                                bc_y: Tuple[BoundaryCondition, BoundaryCondition],
                                n_modes: int = 10) -> Tuple[np.ndarray, np.ndarray]:
        """
        Solve the eigenvalue problem for free vibration.
        
        Args:
            bc_x: Boundary conditions in x-direction (x=0, x=a)
            bc_y: Boundary conditions in y-direction (y=0, y=b)
            n_modes: Number of modes to extract
            
        Returns:
            frequencies: Natural frequencies (rad/s)
            mode_shapes: Mode shapes
        """
        # Assemble matrices
        K, M = self.assemble_matrices()
        
        # Apply boundary conditions
        K_reduced, M_reduced = self.apply_boundary_conditions(K, M, bc_x, bc_y)
        
        # Solve generalized eigenvalue problem
        eigenvals, eigenvecs = eig(K_reduced, M_reduced)
        
        # Sort by frequency
        idx = np.argsort(eigenvals.real)
        eigenvals = eigenvals[idx]
        eigenvecs = eigenvecs[:, idx]
        
        # Convert to frequencies
        frequencies = np.sqrt(eigenvals.real[:n_modes])
        mode_shapes = eigenvecs[:, :n_modes]
        
        return frequencies, mode_shapes
    
    def dimensionless_frequency(self, freq: float, reference_material: MaterialProperties) -> float:
        """
        Calculate dimensionless frequency parameter.
        
        Args:
            freq: Natural frequency (rad/s)
            reference_material: Reference material for normalization
            
        Returns:
            Dimensionless frequency
        """
        return freq * self.plate.length**2 / self.plate.total_thickness * \
               np.sqrt(reference_material.rho / reference_material.E)

# Example usage and validation
if __name__ == "__main__":
    from fg_grc_plate import GrapheneReinforcedComposite, MaterialProperties, GrapheneGeometry, DistributionPattern
    
    # Material properties
    epoxy = MaterialProperties(E=3.0, nu=0.34, rho=1200)
    graphene = MaterialProperties(E=1010, nu=0.186, rho=1060)
    gpl_geometry = GrapheneGeometry()
    
    # Create GRC material
    grc = GrapheneReinforcedComposite(epoxy, graphene, gpl_geometry)
    
    # Test case 1: Pure epoxy plate (validation)
    layers_epoxy = [Layer(thickness=2.0e-3, material=epoxy)]
    plate_epoxy = LaminatedPlate(layers_epoxy, length=0.1, width=0.1)
    
    # Test case 2: Three-layer [0°/GRC/90°] plate
    layers_grc = [
        Layer(thickness=0.5e-3, material=epoxy, angle=0.0),
        Layer(thickness=1.0e-3, material=epoxy, is_grc=True, 
              grc_material=grc, grc_pattern=DistributionPattern.FG_X,
              grc_volume_fraction=0.01, grc_exponent=1.0),
        Layer(thickness=0.5e-3, material=epoxy, angle=90.0)
    ]
    plate_grc = LaminatedPlate(layers_grc, length=0.1, width=0.1)
    
    # Create solver
    solver_epoxy = PlateVibrationSolver(plate_epoxy, n_points_x=9, n_points_y=9)
    solver_grc = PlateVibrationSolver(plate_grc, n_points_x=9, n_points_y=9)
    
    # Solve for simply supported boundary conditions
    bc_ss = (BoundaryCondition.SIMPLY_SUPPORTED, BoundaryCondition.SIMPLY_SUPPORTED)
    
    print("Vibration Analysis Results:")
    print("=" * 50)
    
    # Pure epoxy plate
    freqs_epoxy, modes_epoxy = solver_epoxy.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=5)
    print("\nPure Epoxy Plate (SSSS):")
    for i, freq in enumerate(freqs_epoxy):
        dim_freq = solver_epoxy.dimensionless_frequency(freq, epoxy)
        print(f"Mode {i+1}: {freq:.2f} rad/s, Dimensionless: {dim_freq:.4f}")
    
    # GRC plate
    freqs_grc, modes_grc = solver_grc.solve_eigenvalue_problem(bc_ss, bc_ss, n_modes=5)
    print(f"\nGRC Plate [0°/GRC/90°] with 1% GPL (SSSS):")
    for i, freq in enumerate(freqs_grc):
        dim_freq = solver_grc.dimensionless_frequency(freq, epoxy)
        enhancement = freq / freqs_epoxy[i] if i < len(freqs_epoxy) else 1.0
        print(f"Mode {i+1}: {freq:.2f} rad/s, Dimensionless: {dim_freq:.4f}, Enhancement: {enhancement:.2f}")
    
    print(f"\nOverall frequency enhancement: {freqs_grc[0]/freqs_epoxy[0]:.2f}")
    print("Vibration analysis implementation completed successfully!")